import argparse
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from dataset.CamVid import CamVid
import os
from model.build_BiSeNet import BiSeNet
import torch
from tensorboardX import SummaryWriter
import tqdm
import numpy as np
from utils import poly_lr_scheduler
from utils import reverse_one_hot, compute_global_accuracy, fast_hist, \
    per_class_iu
from loss import DiceLoss


def val(args, model, dataloader):
    print('start val!')
    # label_info = get_label_info(csv_path)
    with torch.no_grad():
        model.eval()
        precision_record = []
        hist = np.zeros((args.num_classes, args.num_classes))
        for i, (data, label) in enumerate(dataloader):
            if torch.cuda.is_available() and args.use_gpu:
                data = data.cuda()
                label = label.cuda()

            # get RGB predict image
            predict = model(data).squeeze()
            predict = reverse_one_hot(predict)
            predict = predict.cpu().numpy()
            print(predict.shape,label.size)
            # get RGB label image
            label = label.squeeze()
            if args.loss == 'dice':
                label = reverse_one_hot(label)
            label = label.cpu().numpy()

            # compute per pixel accuracy

            precision = compute_global_accuracy(predict, label)
            hist += fast_hist(label.flatten(),
                              predict.flatten(), args.num_classes)

            # there is no need to transform the one-hot array to visual RGB array
            # predict = colour_code_segmentation(np.array(predict), label_info)
            # label = colour_code_segmentation(np.array(label), label_info)
            precision_record.append(precision)
            
        precision = np.mean(precision_record)
        # miou = np.mean(per_class_iu(hist))
        miou_list = per_class_iu(hist)[:-1]
        # miou_dict, miou = cal_miou(miou_list, csv_path)
        miou = np.mean(miou_list)
        mem = f'   {torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G'  # (GB)
        print('mIoU for validation: %.3f GPU:%s pixel:%.3f' % (miou,mem,precision))
        # miou_str = ''
        # for key in miou_dict:
        #     miou_str += '{}:{},\n'.format(key, miou_dict[key])
        # print('mIoU for each class:')
        # print(miou_str)
        return precision, miou


def train(args, model, optimizer, dataloader_train, dataloader_val):
    writer = SummaryWriter(comment=''.format(
        args.optimizer, args.context_path))
    if args.loss == 'dice':
        loss_func = DiceLoss()
    elif args.loss == 'crossentropy':
        loss_func = torch.nn.CrossEntropyLoss()
    max_miou = 0
    step = 0
    last_loss_train_mean,min_loss_train_mean=10000,10000
    for epoch in range(args.num_epochs):
        lr = poly_lr_scheduler(optimizer, args.learning_rate,
                               iter=epoch, max_iter=args.num_epochs)
        model.train()
        # tq = tqdm.tqdm(total=len(dataloader_train) * args.batch_size)
        print('epoch %d, lr %f' % (epoch, lr))
        loss_record = []
        for i, (data, label) in enumerate(dataloader_train):
            if torch.cuda.is_available() and args.use_gpu:
                data = data.cuda()
                label = label.cuda()
            output, output_sup1, output_sup2 = model(data)
            loss1 = loss_func(output, label)
            loss2 = loss_func(output_sup1, label)
            loss3 = loss_func(output_sup2, label)
            loss = loss1 + loss2 + loss3
            # tq.update(args.batch_size)
            # tq.set_postfix(loss='%.6f' % loss)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            step += 1
            writer.add_scalar('loss_step', loss, step)
            loss_record.append(loss.item())
        # tq.close()
        loss_train_mean = np.mean(loss_record)
        writer.add_scalar('epoch/loss_epoch_train',
                          float(loss_train_mean), epoch)
            
        print('loss for train : %f' % (loss_train_mean))
        if epoch % args.checkpoint_step == 0 and epoch != 0 and loss_train_mean<min_loss_train_mean:
            print('save loss for train : %f' % (loss_train_mean))
            min_loss_train_mean=loss_train_mean
            if not os.path.isdir(args.save_model_path):
                os.mkdir(args.save_model_path)
            torch.save(model.module.state_dict(),
                       os.path.join(args.save_model_path, 'latest_dice_loss.pth'))

        if epoch % args.validation_step == 0:
            precision, miou = val(args, model, dataloader_val)
            if miou > max_miou:
                max_miou = miou
                torch.save(model.module.state_dict(),
                           os.path.join(args.save_model_path, 'best_dice_loss.pth'))
            writer.add_scalar('epoch/precision_val', precision, epoch)
            writer.add_scalar('epoch/miou val', miou, epoch)


def main(params):
    # basic parameters
    parser = argparse.ArgumentParser()
    parser.add_argument('--num_epochs', type=int, default=300,
                        help='Number of epochs to train for')
    parser.add_argument('--epoch_start_i', type=int, default=0,
                        help='Start counting epochs from this number')
    parser.add_argument('--checkpoint_step', type=int, default=1,
                        help='How often to save checkpoints (epochs)')
    parser.add_argument('--validation_step', type=int, default=1,
                        help='How often to perform validation (epochs)')
    parser.add_argument('--dataset', type=str,
                        default="CamVid", help='Dataset you are using.')
    parser.add_argument('--crop_height', type=int, default=720,
                        help='Height of cropped/resized input image to network')
    parser.add_argument('--crop_width', type=int, default=960,
                        help='Width of cropped/resized input image to network')
    parser.add_argument('--batch_size', type=int, default=1,
                        help='Number of images in each batch')
    parser.add_argument('--context_path', type=str, default="resnet101",
                        help='The context path model you are using, resnet18, resnet101.')
    parser.add_argument('--learning_rate', type=float,
                        default=0.01, help='learning rate used for train')
    parser.add_argument('--data', type=str, default='',
                        help='path of training data')
    parser.add_argument('--num_workers', type=int,
                        default=4, help='num of workers')
    parser.add_argument('--num_classes', type=int, default=32,
                        help='num of object classes (with void)')
    parser.add_argument('--cuda', type=str, default='0',
                        help='GPU ids used for training')
    parser.add_argument('--use_gpu', type=bool, default=True,
                        help='whether to user gpu for training')
    parser.add_argument('--pretrained_model_path', type=str,
                        default=None, help='path to pretrained model')
    parser.add_argument('--save_model_path', type=str,
                        default=None, help='path to save model')
    parser.add_argument('--optimizer', type=str, default='rmsprop',
                        help='optimizer, support rmsprop, sgd, adam')
    parser.add_argument('--loss', type=str, default='dice',
                        help='loss function, dice or crossentropy')
    parser.add_argument('--csv', type=str, default='dice',
                        help='loss function, dice or crossentropy')
    parser.add_argument('--train_label_path', type=str, default='dice',
                        help='loss function, dice or crossentropy')
    parser.add_argument('--test_label_path', type=str, default='dice',
                        help='loss function, dice or crossentropy')

    args = parser.parse_args(params)
    print(args)
    # create dataset and dataloader
    train_path = [os.path.join(args.data, 'train'),
                  os.path.join(args.data, 'val')]
    test_path = os.path.join(args.data, 'test')
    dataset_train = CamVid(train_path, args.train_label_path, args.csv, scale=(args.crop_height, args.crop_width),
                           loss=args.loss, mode='train')
    dataloader_train = DataLoader(
        dataset_train,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
        drop_last=True
    )
    dataset_val = CamVid(test_path, args.test_label_path,  args.csv, scale=(args.crop_height, args.crop_width),
                         loss=args.loss, mode='test')
    dataloader_val = DataLoader(
        dataset_val,
        # this has to be 1
        batch_size=1,
        shuffle=True,
        num_workers=args.num_workers
    )

    # build model
    os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
    model = BiSeNet(args.num_classes, args.context_path)
    if torch.cuda.is_available() and args.use_gpu:
        model = torch.nn.DataParallel(model).cuda()

    # build optimizer
    if args.optimizer == 'rmsprop':
        optimizer = torch.optim.RMSprop(model.parameters(), args.learning_rate)
    elif args.optimizer == 'sgd':
        optimizer = torch.optim.SGD(
            model.parameters(), args.learning_rate, momentum=0.9, weight_decay=1e-4)
    elif args.optimizer == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), args.learning_rate)
    else:  # rmsprop
        print('not supported optimizer \n')
        return None

    # load pretrained model if exists
    # if args.pretrained_model_path is not None:
    #     print('load model from %s ...' % args.pretrained_model_path)
    #     model.module.load_state_dict(torch.load(args.pretrained_model_path))
    #     print('Done!')

    # train
    train(args, model, optimizer, dataloader_train, dataloader_val)

    # val(args, model, dataloader_val, csv_path)

import sys 
if __name__ == '__main__':
    if str(sys.argv[1])=="zhatu":
        print("zhatu")
        params_luotu = [
        '--num_epochs', '1000',
        '--learning_rate', '2.5e-2',
        '--data', '/path/4/CamVid',
        '--num_workers', '1',
        '--num_classes', '4',
        '--cuda', '0',
        # '--pretrained_model_path','/project/train/models/best_dice_loss.pth',
        '--batch_size', '2',  # 6 for resnet101, 12 for resnet18
        '--save_model_path', '/project/train/models',
        '--context_path', 'resnet18',  # only support resnet18 and resnet101
        '--optimizer', 'sgd',
        '--train_label_path','/home/data/1441/*.jpg',
        '--test_label_path','/home/data/1441/*.jpg',
        '--csv', 'zhatu.txt'
    ]
    elif  str(sys.argv[1])=="jingai":
        print("jingai")
        params_luotu = [
        '--num_epochs', '1000',
        '--learning_rate', '2.5e-2',
        '--data', '/path/4/CamVid',
        '--num_workers', '1',
        '--num_classes', '3',
        '--cuda', '0',
        # '--pretrained_model_path','/project/train/models/best_dice_loss.pth',
        '--batch_size', '5',  # 6 for resnet101, 12 for resnet18
        '--save_model_path', '/project/train/models',
        '--context_path', 'resnet18',  # only support resnet18 and resnet101
        '--optimizer', 'sgd',
        '--train_label_path','/home/data/1430/*.jpg',
        '--test_label_path','/home/data/1433/*.jpg',
        '--csv', 'jingai.txt'
    ]
    elif  str(sys.argv[1])=="water":
        print("water")
        params_luotu = [
        '--num_epochs', '1000',
        '--learning_rate', '2.5e-2',
        '--data', '/path/4/CamVid',
        '--num_workers', '1',
        '--num_classes', '2',
        '--cuda', '0',
        # '--pretrained_model_path','/project/train/models/best_dice_loss.pth',
        '--batch_size', '5',  # 6 for resnet101, 12 for resnet18
        '--save_model_path', '/project/train/models',
        '--context_path', 'resnet18',  # only support resnet18 and resnet101
        '--optimizer', 'sgd',
        '--train_label_path','/home/data/572/*.jpg',
        '--test_label_path','/home/data/573/*.jpg',
        '--csv', 'jishui.txt'
    ]
    elif str(sys.argv[1])=="luotu":
        params_luotu = [
            '--num_epochs', '1000',
            '--learning_rate', '2.5e-2',
            '--data', '/path/to/CamVid',
            '--num_workers', '2',
            '--num_classes', '2',
            '--cuda', '0',
            # '--pretrained_model_path','/project/train/models/latest_dice_loss.pth',
            '--batch_size', '12',  # 6 for resnet101, 12 for resnet18
            '--save_model_path', '/project/train/models',
            '--context_path', 'resnet18',  # only support resnet18 and resnet101
            '--optimizer', 'sgd',
            '--csv', 'luotu.csv',
            '--train_label_path','/home/data/865/*.jpg',
            '--test_label_path','/home/data/865/*.jpg',
            '--validation_step','1000',
        ]
    elif str(sys.argv[1])=="truck":
        params_luotu = [
            '--num_epochs', '1000',
            '--learning_rate', '2.5e-2',
            '--data', '/path/to/CamVid',
            '--num_workers', '2',
            '--num_classes', '4',
            '--cuda', '0',
            '--pretrained_model_path','/project/train/models/latest_dice_loss.pth', #'/project/train/models/latest_dice_loss.pth',
            '--batch_size', '10',  # 6 for resnet101, 12 for resnet18
            '--save_model_path', '/project/train/models',
            '--context_path', 'resnet18',  # only support resnet18 and resnet101
            '--optimizer', 'sgd',
            '--csv', 'zhatu.txt',
            '--train_label_path','/home/data/truck/*.jpg',
            '--test_label_path','/home/data/truck/*.jpg',
            '--validation_step','3',
        ]
    elif str(sys.argv[1])=="deng":
        params_luotu = [
            '--num_epochs', '10',
            '--learning_rate', '2.5e-2',
            '--data', '/path/to/CamVid',
            '--num_workers', '2',
            '--num_classes', '4',
            '--cuda', '0',
            # '--pretrained_model_path','/project/train/models/latest_dice_loss.pth', #'/project/train/models/latest_dice_loss.pth',
            '--batch_size', '10',  # 6 for resnet101, 12 for resnet18
            '--save_model_path', '/project/train/models',
            '--context_path', 'resnet18',  # only support resnet18 and resnet101
            '--optimizer', 'sgd',
            '--csv', 'zhatu.txt',
            '--train_label_path','/home/data/1592/*.jpg',
            '--test_label_path','/home/data/1592/*.jpg',
            '--validation_step','3',
        ]
    main(params_luotu)
